#ifndef __MULTIGRID_H__
#define __MULTIGRID_H__

#include "Operator.h"
#include <ostream>
#include <fstream>
#include <utility>
#include <iomanip>
#include <valarray>

#define TEMPLATE template <typename Element, template <typename> class Restriction, template <typename> class Interpolation>
#define Template template <typename Element, template <typename, template <typename> class, template <typename> class> class Cycle, template <typename> class Restriction, template <typename> class Interpolation>

template
<
	typename Element = P1Element,
	template <typename> class Restriction = FullWeighting,
	template <typename> class Interpolation = Linear
>
class VCycle
{
protected:
	size_t n;
	size_t v1;
	size_t v2;
	std::valarray<Point<2> > domain;
	const Vector SolveLayer(const Vector& initial_value, const SpMat& A, const Vector& f_h, size_t layer_n) const;

public:
	VCycle() = default;
	VCycle(size_t _n, const std::valarray<Point<2> >& _domain, size_t _v1 = 2, size_t _v2 = 2);
	virtual ~VCycle() = default;
	void SetN(size_t _n);
	size_t GetN() const;
	void SetV(size_t _v1, size_t _v2);
	const std::valarray<Point<2> >& GetDomain() const;
	virtual const Vector Solve(const Vector& initial_value) const;
};

TEMPLATE
VCycle<Element, Restriction, Interpolation>::VCycle(size_t _n, const std::valarray<Point<2> >& _domain, size_t _v1, size_t _v2) :
	n(_n), domain(_domain), v1(_v1), v2(_v2) {}

TEMPLATE
void VCycle<Element, Restriction, Interpolation>::SetN(size_t _n)
{
	n = _n;
}

TEMPLATE
size_t VCycle<Element, Restriction, Interpolation>::GetN() const
{
	return n;
}

TEMPLATE
void VCycle<Element, Restriction, Interpolation>::SetV(size_t _v1, size_t _v2)
{
	v1 = _v1;
	v2 = _v2;
}

TEMPLATE
const std::valarray<Point<2> >& VCycle<Element, Restriction, Interpolation>::GetDomain() const
{
	return domain;
}

TEMPLATE
const Vector VCycle<Element, Restriction, Interpolation>::Solve(const Vector& initial_value) const
{
	Element element;
	if (initial_value.rows() != element.GetNodeNum(n))
		throw std::length_error("VCycle.Solve(): initial_value.size error!");
	Poisson<Element> poisson(n, domain);
	return SolveLayer(initial_value, poisson.GetMatrix(), poisson.GetRhs(), n);
}

TEMPLATE
const Vector VCycle<Element, Restriction, Interpolation>::SolveLayer(const Vector& initial_value, const SpMat& A, const Vector& f_h, size_t layer_n) const
{
	Element element;
	if (initial_value.rows() != element.GetNodeNum(layer_n) || A.rows() != element.GetNodeNum(layer_n))
		throw std::length_error("VCycle.SolveLayer(): initial_value.size or A.size error!");
	Vector v_h(initial_value), v_2h, f_2h;
	Eigen::ConjugateGradient<Eigen::SparseMatrix<double> > solver;
	solver.compute(A);
	solver.setTolerance(1e-20);

	solver.setMaxIterations(v1);
	v_h = solver.solveWithGuess(f_h, v_h);
	if (layer_n >= 4) {
		Restriction<Element> restriction(layer_n);
		Interpolation<Element> interpolation(layer_n);
		f_2h = restriction * (f_h - A * v_h);
		v_2h = Vector(element.GetNodeNum(layer_n / 2));
		v_2h = SolveLayer(v_2h, Poisson<Element>(layer_n / 2, domain).GetMatrix(), f_2h, layer_n / 2);
		v_h += interpolation * v_2h;
	}
	solver.setMaxIterations(v2);
	v_h = solver.solveWithGuess(f_h, v_h);

	return v_h;
}



template
<
	typename Element = P1Element,
	template <typename> class Restriction = FullWeighting,
	template <typename> class Interpolation = Linear
>
class FullMultigridVCycle : public VCycle<Element, Restriction, Interpolation>
{
private:
	const Vector SolveLayer(const SpMat& A, const Vector& f_h, size_t layer_n) const;

public:
	FullMultigridVCycle() = default;
	FullMultigridVCycle(size_t _n, const std::valarray<Point<2> >& _domain, size_t _v1 = 2, size_t _v2 = 2);
	~FullMultigridVCycle() = default;
	virtual const Vector Solve(const Vector& initial_value = Vector(1)) const;
};

TEMPLATE
FullMultigridVCycle<Element, Restriction, Interpolation>::FullMultigridVCycle(size_t _n, const std::valarray<Point<2> >& _domain, size_t _v1, size_t _v2) : VCycle<Element, Restriction, Interpolation>(_n, _domain, _v1, _v2) {} 


TEMPLATE
const Vector FullMultigridVCycle<Element, Restriction, Interpolation>::Solve(const Vector& initial_value) const
{
	Element element;
	Poisson<Element> poisson(VCycle<Element, Restriction, Interpolation>::n, 
							VCycle<Element, Restriction, Interpolation>::domain);
	return SolveLayer(poisson.GetMatrix(), poisson.GetRhs(), VCycle<Element, Restriction, Interpolation>::n);
}

TEMPLATE
const Vector FullMultigridVCycle<Element, Restriction, Interpolation>::SolveLayer(const SpMat& A, const Vector& f_h, size_t layer_n) const
{
	Element element;
	Poisson<Element> poisson(layer_n, VCycle<Element, Restriction, Interpolation>::domain);
	Vector v_h, v_2h, f_2h;

	if (layer_n >= 4) {
		Restriction<Element> restriction(layer_n);
		Interpolation<Element> interpolation(layer_n);
		f_2h = restriction * f_h;
		v_2h = SolveLayer(Poisson<Element>(layer_n / 2, VCycle<Element, Restriction, Interpolation>::domain).GetMatrix(), f_2h, layer_n / 2);
		v_h = interpolation * v_2h;
	}
	else {
		v_h = Vector(element.GetNodeNum(layer_n));
		v_h.setZero(element.GetNodeNum(layer_n));
	}
	for (size_t i = 0; i < 5; ++i)
		v_h = VCycle<Element, Restriction, Interpolation>::SolveLayer(v_h, A, f_h, layer_n);
	
	return v_h;
}


// to help factory
class AbstractMultigrid
{
public:
	AbstractMultigrid() = default;
	virtual ~AbstractMultigrid() = default;
	virtual void SetParams(size_t _n, const Vector& _init_guess, const std::valarray<Point<2> >& _domain, size_t _max_iterations = 100, double tolerance = 1e-16, size_t _v1 = 2, size_t _v2 = 2) = 0;
	virtual void SetInitialGuess(const Vector& v) = 0;
	virtual void SetMaxIterations(size_t _max_iterations) = 0;
	virtual void SetTolerance(double _tolerance) = 0;
	virtual const Vector Solve() const = 0;
};

template
<
	typename Element = P1Element,
	template <typename, template <typename> class, template <typename> class> class Cycle = VCycle,
	template <typename> class Restriction = FullWeighting,
	template <typename> class Interpolation = Linear
>
class Multigrid : public AbstractMultigrid
{
private:
	Cycle<Element, Restriction, Interpolation> cycle;
	size_t max_iterations;
	double tolerance;
	Vector init_guess;

public:
	Multigrid() = default;
	Multigrid(size_t _n, const Vector& _init_guess, const std::valarray<Point<2> >& _domain, size_t _max_iterations = 100, double tolerance = 1e-16, size_t _v1 = 2, size_t _v2 = 2);
	~Multigrid() = default;
	void SetParams(size_t _n, const Vector& _init_guess, const std::valarray<Point<2> >& _domain, size_t _max_iterations = 100, double _tolerance = 1e-16, size_t _v1 = 2, size_t _v2 = 2);
	void SetInitialGuess(const Vector& v);
	void SetMaxIterations(size_t _max_iterations);
	void SetTolerance(double _tolerance);
	void SetCycle(const Cycle<Element, Restriction, Interpolation>& _cycle);
	const Vector Solve() const;
};

Template
Multigrid<Element, Cycle, Restriction, Interpolation>::Multigrid(size_t _n, const Vector& _init_guess, const std::valarray<Point<2> >& _domain, size_t _max_iterations, double _tolerance, size_t _v1, size_t _v2) :
	cycle(_n, _domain, _v1, _v2), max_iterations(_max_iterations),
	tolerance(_tolerance), init_guess(_init_guess) {}

Template
void Multigrid<Element, Cycle, Restriction, Interpolation>::SetParams(size_t _n, const Vector& _init_guess, const std::valarray<Point<2> >& _domain, size_t _max_iterations, double _tolerance, size_t _v1, size_t _v2)
{
	cycle = Cycle<Element, Restriction, Interpolation>(_n, _domain, _v1, _v2);
	max_iterations = _max_iterations;
	tolerance = _tolerance;
	init_guess = _init_guess;
}

Template
void Multigrid<Element, Cycle, Restriction, Interpolation>::SetInitialGuess(const Vector& v)
{
	init_guess = v;
}

Template
void Multigrid<Element, Cycle, Restriction, Interpolation>::SetMaxIterations(size_t _max_iterations)
{
	max_iterations = _max_iterations;
}

Template
void Multigrid<Element, Cycle, Restriction, Interpolation>::SetTolerance(double _tolerance)
{
	tolerance = _tolerance;
}

Template
void Multigrid<Element, Cycle, Restriction, Interpolation>::SetCycle(const Cycle<Element, Restriction, Interpolation>& _cycle)
{
	cycle = _cycle;
}

Template
const Vector Multigrid<Element, Cycle, Restriction, Interpolation>::Solve() const
{
	size_t n = cycle.GetN();
	Vector v(init_guess);
	Poisson<Element> poisson(n, cycle.GetDomain());
	for (size_t i = 0; i < max_iterations; ++i) {
		v =  cycle.Solve(v);
		//std::cout << "iter " << i << ": " << poisson.L2Error(v) << std::endl;
		if (poisson.L2Error(v) < tolerance)
			break;
	}

	return v;
}

#undef TEMPLATE
#undef Template
#endif
